package com.xiashitech.agent.instrumentation.async.inst;

import com.xiashitech.agent.instrumentation.async.boot.AsyncCallback;
import com.xiashitech.agent.instrumentation.async.boot.AsyncFutureHolder;
import com.xiashitech.agent.instrumentation.base.InstPreparedArgs;
import com.xiashitech.agent.instrumentation.base.inst.CommonInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.util.concurrent.CompletableFuture;

import static net.bytebuddy.matcher.ElementMatchers.*;

/**
 * 异步调用的响应插桩
 * @author zhaoyk
 */
public class XSCompletableFutureInstrumentation extends CommonInstrumentation {

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return is(CompletableFuture.class);
    }

    @Override
    public void transform(TypeTransformer typeTransformer) {
        transformMethod(isMethod()
                        .and(named("get"))
                        .and(isPublic()),
                "GetAdvice", typeTransformer);
    }

    public static class GetAdvice {

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.This CompletableFuture<?> completableFuture,
                @Advice.Return Object value) {

            //DebugUtils.print4Debug("XSCompletableFutureInstrumentation.onExit", completableFuture.getClass());

            //当前织入类CompletableFuture，类加载器bootstrap，所以AsyncFutureHolder需要放在boot包下
            AsyncCallback<?, ?> callback = AsyncFutureHolder.fetch(completableFuture);
            if(callback == null)
                return;

            InstPreparedArgs preparedArgs = InstPreparedArgs.prepareOnEnter("XSCompletableFutureInstrumentation.get", true);
            if(!InstPreparedArgs.checkOnEnter(preparedArgs))
                return;

            //DebugUtils.print4Debug("XSCompletableFutureInstrumentation.complete callback", callback);
            //DebugUtils.print4Debug("XSCompletableFutureInstrumentation.complete rtn", value);

            callback.callback(value, preparedArgs.getContext());
        }
    }

}
